# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
FSDP PPO Trainer with Ray-based single controller.
This trainer supports model-agonistic model initialization with huggingface
"""

import os
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from enum import Enum
from pprint import pprint
from typing import Type, Dict
from collections import defaultdict
import json

import numpy as np
from codetiming import Timer
from omegaconf import OmegaConf, open_dict
from verl import DataProto
from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto, DataProtoItem
from verl.single_controller.base import Worker
from verl.single_controller.ray import RayResourcePool, RayWorkerGroup, RayClassWithInitArgs
from verl.single_controller.ray.base import create_colocated_worker_cls
from verl.trainer.ppo import core_algos
from verl.utils.seqlen_balancing import get_seqlen_balanced_partitions, log_seqlen_unbalance
from deepscaler.globals import THOUGHT_DELIMITER_START, THOUGHT_DELIMITER_END, OAI_RM_MODEL, NO_THINK_PROMPT, NO_THINK_PROMPT_IDS
from torch.nn.utils.rnn import pad_sequence

WorkerType = Type[Worker]

from torch.nn.utils.rnn import pad_sequence
import torch
import numpy as np

def append_no_think_prompt(
    batch: "DataProto",
    no_think_ids: list[int],
    probability: float = 0.1,
    pad_token_id: int = None
) -> "DataProto":
    """
    For each batch:
    - With probability, append `no_think_ids` to all samples.
    - Otherwise, do nothing.

    Args:
        batch (DataProto): Input batch.
        no_think_ids (list[int]): Token ids to append if selected.
        probability (float): Probability to modify the entire batch.
        pad_token_id (int, optional): Padding token id (default to 0).

    Returns:
        DataProto: Modified batch (in-place).
    """
    assert 'input_ids' in batch.batch and 'attention_mask' in batch.batch and 'position_ids' in batch.batch, \
        "Batch must contain input_ids, attention_mask, and position_ids."

    if np.random.rand() >= probability:
        # 大概率直接返回原batch（什么都不动）
        return batch

    input_ids_list = []
    attention_mask_list = []
    position_ids_list = []

    no_think_tensor = torch.tensor(no_think_ids, dtype=batch.batch['input_ids'].dtype, device=batch.batch['input_ids'].device)

    for i in range(len(batch.batch['input_ids'])):
        input_ids = batch.batch['input_ids'][i]
        attention_mask = batch.batch['attention_mask'][i]
        position_ids = batch.batch['position_ids'][i]

        # ====== 仍然使用你原版的追加逻辑 ======
        input_ids = torch.cat([input_ids, no_think_tensor], dim=0)
        attention_mask = torch.cat([attention_mask, torch.ones_like(no_think_tensor)], dim=0)
        position_ids = torch.cat(
            [position_ids, position_ids[-1:] + torch.arange(1, len(no_think_tensor) + 1, device=position_ids.device)],
            dim=0
        )

        input_ids_list.append(input_ids)
        attention_mask_list.append(attention_mask)
        position_ids_list.append(position_ids)

    pad_value = pad_token_id if pad_token_id is not None else 0

    batch.batch['input_ids'] = pad_sequence(input_ids_list, batch_first=True, padding_value=pad_value)
    batch.batch['attention_mask'] = pad_sequence(attention_mask_list, batch_first=True, padding_value=0)
    batch.batch['position_ids'] = pad_sequence(position_ids_list, batch_first=True, padding_value=0)

    return batch




def dataprotoitem_to_dataproto(item: DataProtoItem) -> DataProto:
    """Convert a DataProtoItem to a DataProto object"""
    return DataProto.from_dict(
        tensors=item.batch,  # TensorDict is already in correct format
        non_tensors=item.non_tensor_batch,  # Dict is already in correct format 
        meta_info=item.meta_info
    )

class Role(Enum):
    """
    To create more roles dynamically, you can subclass Role and add new members
    """
    Actor = 0
    Rollout = 1
    ActorRollout = 2
    Critic = 3
    RefPolicy = 4
    RewardModel = 5
    ActorRolloutRef = 6


@dataclass
class ResourcePoolManager:
    """
    Define a resource pool specification. Resource pool will be initialized first.
    Mapping
    """
    resource_pool_spec: dict[str, list[int]]
    mapping: dict[Role, str]
    resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict)

    def create_resource_pool(self):
        for resource_pool_name, process_on_nodes in self.resource_pool_spec.items():
            # max_colocate_count means the number of WorkerGroups (i.e. processes) in each RayResourcePool
            # For FSDP backend, we recommend using max_colocate_count=1 that merge all WorkerGroups into one.
            # For Megatron backend, we recommend using max_colocate_count>1 that can utilize different WorkerGroup for differnt models
            resource_pool = RayResourcePool(process_on_nodes=process_on_nodes,
                                            use_gpu=True,
                                            max_colocate_count=1,
                                            name_prefix=resource_pool_name)
            self.resource_pool_dict[resource_pool_name] = resource_pool

    def get_resource_pool(self, role: Role) -> RayResourcePool:
        """Get the resource pool of the worker_cls"""
        return self.resource_pool_dict[self.mapping[role]]


import torch
from verl.utils.torch_functional import masked_mean


def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty='kl'):
    responses = data.batch['responses']
    response_length = responses.size(1)
    token_level_scores = data.batch['token_level_scores']
    batch_size = data.batch.batch_size[0]
    attention_mask = data.batch['attention_mask']
    response_mask = attention_mask[:, -response_length:]

    # compute kl between ref_policy and current policy
    if 'ref_log_prob' in data.batch.keys():
        kld = core_algos.kl_penalty(data.batch['old_log_probs'], data.batch['ref_log_prob'],
                                    kl_penalty=kl_penalty)  # (batch_size, response_length)
        kld = kld * response_mask
        beta = kl_ctrl.value
    else:
        beta = 0
        kld = torch.zeros_like(response_mask, dtype=torch.float32)

    token_level_rewards = token_level_scores - beta * kld

    current_kl = masked_mean(kld, mask=response_mask, axis=-1)  # average over sequence
    current_kl = torch.mean(current_kl, dim=0).item()

    # according to https://github.com/huggingface/trl/blob/951ca1841f29114b969b57b26c7d3e80a39f75a0/trl/trainer/ppo_trainer.py#L837
    kl_ctrl.update(current_kl=current_kl, n_steps=batch_size)
    data.batch['token_level_rewards'] = token_level_rewards

    metrics = {'critic/kl': current_kl, 'critic/kl_coeff': beta}

    return data, metrics


def compute_advantage(data: DataProto, adv_estimator, gamma=1.0, lam=1.0, num_repeat=1):
    # prepare response group
    # TODO: add other ways to estimate advantages
    if adv_estimator == 'gae':
        values = data.batch['values']
        responses = data.batch['responses']
        response_length = responses.size(-1)
        attention_mask = data.batch['attention_mask']
        response_mask = attention_mask[:, -response_length:]
        token_level_rewards = data.batch['token_level_rewards']
        advantages, returns = core_algos.compute_gae_advantage_return(token_level_rewards=token_level_rewards,
                                                                      values=values,
                                                                      eos_mask=response_mask,
                                                                      gamma=gamma,
                                                                      lam=lam)
        data.batch['advantages'] = advantages
        data.batch['returns'] = returns
    elif adv_estimator == 'grpo':
        token_level_rewards = data.batch['token_level_rewards']
        index = data.non_tensor_batch['uid']
        responses = data.batch['responses']
        response_length = responses.size(-1)
        attention_mask = data.batch['attention_mask']
        response_mask = attention_mask[:, -response_length:]
        advantages, returns = core_algos.compute_grpo_outcome_advantage(token_level_rewards=token_level_rewards,
                                                                        eos_mask=response_mask,
                                                                        index=index)
        data.batch['advantages'] = advantages
        data.batch['returns'] = returns
    else:
        raise NotImplementedError
    return data


def reduce_metrics(metrics: dict):
    for key, val in metrics.items():
        metrics[key] = np.mean(val)
    return metrics


def _compute_response_info(batch):
    response_length = batch.batch['responses'].shape[-1]

    prompt_mask = batch.batch['attention_mask'][:, :-response_length]
    response_mask = batch.batch['attention_mask'][:, -response_length:]

    prompt_length = prompt_mask.sum(-1).float()
    response_length = response_mask.sum(-1).float()  # (batch_size,)

    return dict(
        response_mask=response_mask,
        prompt_length=prompt_length,
        response_length=response_length,
    )


def compute_data_metrics(batch, use_critic=True, tokenizer=None):
    # TODO: add response length
    sequence_score = batch.batch['token_level_scores'].sum(-1)
    sequence_reward = batch.batch['token_level_rewards'].sum(-1)

    advantages = batch.batch['advantages']
    returns = batch.batch['returns']

    max_response_length = batch.batch['responses'].shape[-1]

    prompt_mask = batch.batch['attention_mask'][:, :-max_response_length].bool()
    response_mask = batch.batch['attention_mask'][:, -max_response_length:].bool()

    max_prompt_length = prompt_mask.size(-1)

    response_info = _compute_response_info(batch)
    prompt_length = response_info['prompt_length']
    response_length = response_info['response_length']

    valid_adv = torch.masked_select(advantages, response_mask)
    valid_returns = torch.masked_select(returns, response_mask)

    if use_critic:
        values = batch.batch['values']
        valid_values = torch.masked_select(values, response_mask)
        return_diff_var = torch.var(valid_returns - valid_values)
        return_var = torch.var(valid_returns)

    # ==== 按 think/no-think 划分 ====
    think_flags = torch.tensor(batch.non_tensor_batch['think_flags'].tolist(), device=sequence_score.device).bool()
    is_correct_flags = torch.tensor(batch.non_tensor_batch['is_correct_flags'].tolist(), device=sequence_score.device).bool()

    total_samples = think_flags.size(0)

    num_think = think_flags.sum()
    num_no_think = total_samples - num_think

    think_ratio = num_think.float() / total_samples if total_samples > 0 else torch.tensor(0.0, device=sequence_score.device)
    no_think_ratio = num_no_think.float() / total_samples if total_samples > 0 else torch.tensor(0.0, device=sequence_score.device)

    # 选取分组索引
    think_indices = think_flags.nonzero(as_tuple=True)[0]
    no_think_indices = (~think_flags).nonzero(as_tuple=True)[0]

    # 平均 response token数
    avg_response_length_think = response_length[think_indices].mean() if num_think > 0 else torch.tensor(0.0, device=sequence_score.device)
    avg_response_length_no_think = response_length[no_think_indices].mean() if num_no_think > 0 else torch.tensor(0.0, device=sequence_score.device)

    # 平均分数（sequence_score）
    avg_score_think = sequence_score[think_indices].mean() if num_think > 0 else torch.tensor(0.0, device=sequence_score.device)
    avg_score_no_think = sequence_score[no_think_indices].mean() if num_no_think > 0 else torch.tensor(0.0, device=sequence_score.device)

    # 平均 acc 数
    acc_think = is_correct_flags[think_indices].float().mean() if num_think > 0 else torch.tensor(0.0, device=sequence_score.device)
    acc_no_think = is_correct_flags[no_think_indices].float().mean() if num_no_think > 0 else torch.tensor(0.0, device=sequence_score.device)

    metrics = {
        # score
        'critic/score/mean': sequence_score.mean().detach().item(),
        'critic/score/max': sequence_score.max().detach().item(),
        'critic/score/min': sequence_score.min().detach().item(),

        # reward
        'critic/rewards/mean': sequence_reward.mean().detach().item(),
        'critic/rewards/max': sequence_reward.max().detach().item(),
        'critic/rewards/min': sequence_reward.min().detach().item(),

        # advantages
        'critic/advantages/mean': valid_adv.mean().detach().item(),
        'critic/advantages/max': valid_adv.max().detach().item(),
        'critic/advantages/min': valid_adv.min().detach().item(),

        # returns
        'critic/returns/mean': valid_returns.mean().detach().item(),
        'critic/returns/max': valid_returns.max().detach().item(),
        'critic/returns/min': valid_returns.min().detach().item(),

        # accuracy
        'critic/accuracy': is_correct_flags.float().mean().detach().item(),

        **({
            'critic/values/mean': valid_values.mean().detach().item(),
            'critic/values/max': valid_values.max().detach().item(),
            'critic/values/min': valid_values.min().detach().item(),
            'critic/vf_explained_var': (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(),
        } if use_critic else {}),

        # response length
        'response_length/mean': response_length.mean().detach().item(),
        'response_length/max': response_length.max().detach().item(),
        'response_length/min': response_length.min().detach().item(),
        'response_length/clip_ratio': torch.eq(response_length, max_response_length).float().mean().detach().item(),

        # prompt length
        'prompt_length/mean': prompt_length.mean().detach().item(),
        'prompt_length/max': prompt_length.max().detach().item(),
        'prompt_length/min': prompt_length.min().detach().item(),
        'prompt_length/clip_ratio': torch.eq(prompt_length, max_prompt_length).float().mean().detach().item(),

        # think/no-think 分组统计
        'think/think_ratio': think_ratio.detach().item(),
        'think/nothink_ratio': no_think_ratio.detach().item(),
        'think/think_avg_response_length': avg_response_length_think.detach().item(),
        'think/nothink_avg_response_length': avg_response_length_no_think.detach().item(),
        'think/think_avg_score': avg_score_think.detach().item(),
        'think/nothink_avg_score': avg_score_no_think.detach().item(),
        'think/think_accuracy': acc_think.detach().item(),
        'think/nothink_accuracy': acc_no_think.detach().item(),

        
    }

    return metrics



def compute_timing_metrics(batch, timing_raw):
    response_info = _compute_response_info(batch)
    num_prompt_tokens = torch.sum(response_info['prompt_length']).item()
    num_response_tokens = torch.sum(response_info['response_length']).item()
    num_overall_tokens = num_prompt_tokens + num_response_tokens

    num_tokens_of_section = {
        'gen': num_response_tokens,
        **{
            name: num_overall_tokens for name in ['ref', 'values', 'adv', 'update_critic', 'update_actor']
        },
    }

    return {
        **{
            f'timing_s/{name}': value for name, value in timing_raw.items()
        },
        **{
            f'timing_per_token_ms/{name}': timing_raw[name] * 1000 / num_tokens_of_section[name] for name in set(num_tokens_of_section.keys(
            )) & set(timing_raw.keys())
        },
    }


@contextmanager
def _timer(name: str, timing_raw: Dict[str, float]):
    with Timer(name=name, logger=None) as timer:
        yield
    timing_raw[name] = timer.last

from torch.utils.data import RandomSampler, SequentialSampler, Sampler
import random
from typing import Iterator, Dict, List, Optional
class DifficultyStratifiedBatchSampler(Sampler[List[int]]):
    def __init__(
        self,
        dataset,
        batch_size: int,
        ratios: Optional[Dict[str, float]] = None,
        seed: int = 42,
        shuffle: bool = True,
        divisible_by: Optional[int] = None  # 添加对 world_size 的对齐支持
    ):
        self.dataset = dataset
        self.batch_size = batch_size
        self.seed = seed
        self.shuffle = shuffle
        self.divisible_by = divisible_by

        # 处理 ratio
        ratios = ratios or {'easy': 0.1, 'medium': 0.8, 'hard': 0.1}
        ratios = {k: v for k, v in ratios.items() if v > 0.0}
        total = sum(ratios.values())
        if total <= 0.0:
            raise ValueError("All sampling ratios are zero.")
        self.ratios = {k: v / total for k, v in ratios.items()}

        # 构建 difficulty 索引池
        self.difficulty_to_indices = defaultdict(list)
        for i, sample in enumerate(dataset):
            difficulty = sample.get("difficulty", "medium")
            self.difficulty_to_indices[difficulty].append(i)

        self.difficulty_keys = list(self.ratios.keys())
        self.steps_per_epoch = self._estimate_steps_per_epoch()

    def _estimate_steps_per_epoch(self) -> int:
        # 根据最少的 difficulty 类别推断 epoch 步数
        min_steps = float("inf")
        for diff, ratio in self.ratios.items():
            n_available = len(self.difficulty_to_indices[diff])
            n_needed_per_batch = max(1, int(self.batch_size * ratio))
            min_steps = min(min_steps, n_available // n_needed_per_batch)
        return min_steps

    def __iter__(self) -> Iterator[List[int]]:
        rng = random.Random(self.seed)

        # 每轮 epoch 开始前 shuffle 一下
        difficulty_pools = {}
        for diff in self.difficulty_keys:
            indices = self.difficulty_to_indices[diff][:]
            if self.shuffle:
                rng.shuffle(indices)
            difficulty_pools[diff] = indices

        # 按比例每 batch 采样
        ptr = {k: 0 for k in self.difficulty_keys}
        for _ in range(self.steps_per_epoch):
            batch_indices = []
            for diff in self.difficulty_keys:
                n = int(self.batch_size * self.ratios[diff])
                pool = difficulty_pools[diff]
                start = ptr[diff]
                end = start + n
                if end > len(pool):  # 若不够就从头再来
                    rng.shuffle(pool)
                    start = 0
                    end = n
                    ptr[diff] = 0
                batch_indices.extend(pool[start:end])
                ptr[diff] += n

            # 对齐 world_size
            if self.divisible_by is not None:
                remainder = len(batch_indices) % self.divisible_by
                if remainder != 0:
                    pad_len = self.divisible_by - remainder
                    pad_pool = sum(difficulty_pools.values(), [])
                    batch_indices += rng.choices(pad_pool, k=pad_len)

            rng.shuffle(batch_indices)
            yield batch_indices

    def __len__(self) -> int:
        return self.steps_per_epoch

    
class RayPPOTrainer(object):
    """
    Note that this trainer runs on the driver process on a single CPU/GPU node.
    """

    # TODO: support each role have individual ray_worker_group_cls,
    # i.e., support different backend of different role
    def __init__(self,
                 config,
                 tokenizer,
                 role_worker_mapping: dict[Role, WorkerType],
                 resource_pool_manager: ResourcePoolManager,
                 ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup,
                 reward_fn=None,
                 val_reward_fn=None):

        # assert torch.cuda.is_available(), 'cuda must be available on driver'

        self.tokenizer = tokenizer
        self.config = config
        self.reward_fn = reward_fn
        self.val_reward_fn = val_reward_fn

        self.hybrid_engine = config.actor_rollout_ref.hybrid_engine
        assert self.hybrid_engine, 'Currently, only support hybrid engine'

        if self.hybrid_engine:
            assert Role.ActorRollout in role_worker_mapping, f'{role_worker_mapping.keys()=}'

        self.role_worker_mapping = role_worker_mapping
        self.resource_pool_manager = resource_pool_manager
        self.use_reference_policy = Role.RefPolicy in role_worker_mapping
        self.use_rm = Role.RewardModel in role_worker_mapping
        self.ray_worker_group_cls = ray_worker_group_cls

        # define KL control
        if self.use_reference_policy:
            if config.algorithm.kl_ctrl.type == 'fixed':
                self.kl_ctrl = core_algos.FixedKLController(kl_coef=config.algorithm.kl_ctrl.kl_coef)
            elif config.algorithm.kl_ctrl.type == 'adaptive':
                assert config.algorithm.kl_ctrl.horizon > 0, f'horizon must be larger than 0. Got {config.critic.kl_ctrl.horizon}'
                self.kl_ctrl = core_algos.AdaptiveKLController(init_kl_coef=config.algorithm.kl_ctrl.kl_coef,
                                                               target_kl=config.algorithm.kl_ctrl.target_kl,
                                                               horizon=config.algorithm.kl_ctrl.horizon)
            else:
                raise NotImplementedError
        else:
            self.kl_ctrl = core_algos.FixedKLController(kl_coef=0.)

        self._create_dataloader()

    def _create_dataloader(self):
        from torch.utils.data import DataLoader
        # TODO: we have to make sure the batch size is divisible by the dp size
        from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn
        self.train_dataset = RLHFDataset(parquet_files=self.config.data.train_files,
                                         tokenizer=self.tokenizer,
                                         prompt_key=self.config.data.prompt_key,
                                         max_prompt_length=self.config.data.max_prompt_length,
                                         filter_prompts=True,
                                         return_raw_chat=self.config.data.get('return_raw_chat', False),
                                         use_difficulty_sampling=self.config.data.get('use_difficulty_sampling', False),
                                         truncation='error')
        train_batch_size = self.config.data.train_batch_size
        self.modified_dataset = [{"index": self.train_dataset[i]["index"], "reward": -1.0, "acc": -1.0} for i in range(len(self.train_dataset)) ]
        
        if self.config.trainer.rejection_sample:
            train_batch_size *= self.config.trainer.rejection_sample_multiplier
            train_batch_size = int(train_batch_size)

        if self.config.data.get("use_difficulty_sampling", False):
            # sampler = DifficultyStratifiedSampler(
            #     data_source = self.train_dataset,
            #     # difficulty_key="difficulty",
            #     ratios={"easy": self.config.data.sample_easy_ratio, "medium": self.config.data.sample_medium_ratio, "hard": self.config.data.sample_hard_ratio},
            #     seed = self.config.data.get('seed', 1),
            #     shuffle = True
            # )
            batch_sampler = DifficultyStratifiedBatchSampler(
                dataset = self.train_dataset,
                batch_size = self.config.data.get('gen_batch_size',self.config.data.train_batch_size),
                ratios={"easy": self.config.data.sample_easy_ratio, "medium": self.config.data.sample_medium_ratio, "hard": self.config.data.sample_hard_ratio},
                seed = self.config.data.get('seed', 1),
                shuffle = True,
                divisible_by = 4 * 8
            )
            from torchdata.stateful_dataloader import StatefulDataLoader
            self.train_dataloader = StatefulDataLoader(dataset=self.train_dataset,
                                                    num_workers=8,
                                                    collate_fn=collate_fn,
                                                    batch_sampler=batch_sampler)
        else:
            self.train_dataloader = DataLoader(dataset=self.train_dataset,
                                            batch_size=train_batch_size,
                                            shuffle=True,
                                            drop_last=True,
                                            collate_fn=collate_fn)

        self.val_dataset = RLHFDataset(parquet_files=self.config.data.val_files,
                                    tokenizer=self.tokenizer,
                                    prompt_key=self.config.data.prompt_key,
                                    max_prompt_length=self.config.data.max_prompt_length,
                                    filter_prompts=True,
                                    return_raw_chat=self.config.data.get('return_raw_chat', False),
                                    use_difficulty_sampling=self.config.data.get('use_difficulty_sampling', False),
                                    truncation='error')
        self.val_dataloader = DataLoader(dataset=self.val_dataset,
                                         batch_size=len(self.val_dataset),
                                         shuffle=True,
                                         drop_last=True,
                                         collate_fn=collate_fn)

        assert len(self.train_dataloader) >= 1
        assert len(self.val_dataloader) >= 1

        print(f'Size of train dataloader: {len(self.train_dataloader)}')
        print(f'Size of val dataloader: {len(self.val_dataloader)}')

        # inject total_training_steps to actor/critic optim_config. This is hacky.
        total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs

        if self.config.trainer.total_training_steps is not None:
            total_training_steps = self.config.trainer.total_training_steps

        self.total_training_steps = total_training_steps
        print(f'Total training steps: {self.total_training_steps}')

        OmegaConf.set_struct(self.config, True)
        with open_dict(self.config):
            self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps
            self.config.critic.optim.total_training_steps = total_training_steps

    def _validate(self):
        reward_tensor_lst = []
        acc_tensor_lst = []
        think_tensor_lst = []
        data_source_lst = []

        for test_data in self.val_dataloader:
            test_batch = DataProto.from_single_dict(test_data)
            # test_batch = test_batch.to('cuda')

            # we only do validation on rule-based rm
            if self.config.reward_model.enable and test_batch[0].non_tensor_batch['reward_model']['style'] == 'model':
                return {}

            n_val_samples = self.config.actor_rollout_ref.rollout.n_val
            test_batch = test_batch.repeat(repeat_times=n_val_samples, interleave=True)
            test_gen_batch = test_batch.pop(['input_ids', 'attention_mask', 'position_ids'])
            test_gen_batch.meta_info = {
                'eos_token_id': self.tokenizer.eos_token_id,
                'pad_token_id': self.tokenizer.pad_token_id,
                'recompute_log_prob': False,
                'do_sample': False,
                'validate': True,
            }

            # pad to be divisible by dp_size
            test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, self.actor_rollout_wg.world_size)
            test_gen_batch_padded.meta_info['val_temperature'] = self.config.actor_rollout_ref.rollout.val_temperature
            test_output_gen_batch_padded = self.actor_rollout_wg.generate_sequences(test_gen_batch_padded)
            # unpad
            test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size)
            print('Validation: Generation end.')

            test_batch = test_batch.union(test_output_gen_batch)

            # evaluate using reward_function
            # for certain reward function (e.g. sandbox), the generation can overlap with reward
            reward_result = self.val_reward_fn(test_batch, return_dict=True)
            reward_tensor = reward_result['reward_tensor']  # (batch_size, seq_len)
            reward_extra_infos_dict = reward_result['reward_extra_info']

            think_flags = reward_extra_infos_dict.get('think_flags', None)
            is_correct_flags = reward_extra_infos_dict.get('is_correct_flags', None)

            if is_correct_flags is not None:
                acc_tensor = torch.tensor(is_correct_flags, dtype=torch.float32)
            else:
                acc_tensor = torch.ones(reward_tensor.size(0), dtype=torch.float32)

            if think_flags is not None:
                think_tensor = torch.tensor(think_flags, dtype=torch.float32)
            else:
                think_tensor = torch.zeros(reward_tensor.size(0), dtype=torch.float32)

            reward_tensor_lst.append(reward_tensor)
            acc_tensor_lst.append(acc_tensor)
            think_tensor_lst.append(think_tensor)
            data_source_lst.append(test_batch.non_tensor_batch.get('data_source', ['unknown'] * reward_tensor.shape[0]))

        reward_tensor = torch.cat(reward_tensor_lst, dim=0).sum(-1).cpu()  # (batch_size,)
        acc_tensor = torch.cat(acc_tensor_lst, dim=0).cpu()
        think_tensor = torch.cat(think_tensor_lst, dim=0).cpu()
        data_sources = np.concatenate(data_source_lst, axis=0)

        # Group by data_source
        data_source_reward = {}
        data_source_acc = {}
        data_source_think = {}

        for i in range(reward_tensor.shape[0]):
            data_source = data_sources[i]
            if data_source not in data_source_reward:
                data_source_reward[data_source] = []
                data_source_acc[data_source] = []
                data_source_think[data_source] = []

            data_source_reward[data_source].append(reward_tensor[i].item())
            data_source_acc[data_source].append(acc_tensor[i].item())
            data_source_think[data_source].append(think_tensor[i].item())

        metric_dict = {}
        for data_source in data_source_reward:
            rewards = data_source_reward[data_source]
            accs = data_source_acc[data_source]
            thinks = data_source_think[data_source]

            metric_dict[f'val/test_score/{data_source}'] = np.mean(rewards)
            metric_dict[f'val/test_acc/{data_source}'] = np.mean(accs)
            metric_dict[f'val/test_think_ratio/{data_source}'] = np.mean(thinks)

        return metric_dict

    def init_workers(self):
        """Init resource pool and worker group"""
        self.resource_pool_manager.create_resource_pool()

        self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}

        # create actor and rollout
        if self.hybrid_engine:
            resource_pool = self.resource_pool_manager.get_resource_pool(Role.ActorRollout)
            actor_rollout_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.ActorRollout],
                                                     config=self.config.actor_rollout_ref,
                                                     role='actor_rollout')
            self.resource_pool_to_cls[resource_pool]['actor_rollout'] = actor_rollout_cls
        else:
            raise NotImplementedError

        # create critic
        if self.config.algorithm.adv_estimator == 'gae':
            resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic)
            critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=self.config.critic)
            self.resource_pool_to_cls[resource_pool]['critic'] = critic_cls
            self.use_critic = True
        elif self.config.algorithm.adv_estimator == 'grpo':
            self.use_critic = False
        else:
            raise NotImplementedError

        # create reference policy if needed
        if self.use_reference_policy:
            resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy)
            ref_policy_cls = RayClassWithInitArgs(self.role_worker_mapping[Role.RefPolicy],
                                                  config=self.config.actor_rollout_ref,
                                                  role='ref')
            self.resource_pool_to_cls[resource_pool]['ref'] = ref_policy_cls

        # create a reward model if reward_fn is None
        if self.use_rm:
            # we create a RM here
            resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel)
            rm_cls = RayClassWithInitArgs(self.role_worker_mapping[Role.RewardModel], config=self.config.reward_model)
            self.resource_pool_to_cls[resource_pool]['rm'] = rm_cls

        # initialize WorkerGroup
        # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,
        # you should not use `create_colocated_worker_cls`. Instead, directly pass different resource pool to different worker groups.
        # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.
        all_wg = {}
        self.wg_dicts = []
        for resource_pool, class_dict in self.resource_pool_to_cls.items():
            worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
            wg_dict = self.ray_worker_group_cls(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls)
            spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
            all_wg.update(spawn_wg)
            # keep the referece of WorkerDict to support ray >= 2.31. Ref: https://github.com/ray-project/ray/pull/45699
            self.wg_dicts.append(wg_dict)

        if self.use_critic:
            self.critic_wg = all_wg['critic']
            self.critic_wg.init_model()

        if self.use_reference_policy:
            self.ref_policy_wg = all_wg['ref']
            self.ref_policy_wg.init_model()

        if self.use_rm:
            self.rm_wg = all_wg['rm']
            self.rm_wg.init_model()

        # we should create rollout at the end so that vllm can have a better estimation of kv cache memory
        self.actor_rollout_wg = all_wg['actor_rollout']
        self.actor_rollout_wg.init_model()

    def _save_checkpoint(self):
        actor_local_path = os.path.join(self.config.trainer.default_local_dir, 'actor',
                                        f'global_step_{self.global_steps}')
        actor_remote_path = None if self.config.trainer.default_hdfs_dir is None else os.path.join(
            self.config.trainer.default_hdfs_dir, 'actor')
        self.actor_rollout_wg.save_checkpoint(actor_local_path, actor_remote_path)

        if self.use_critic:
            critic_local_path = os.path.join(self.config.trainer.default_local_dir, 'critic',
                                             f'global_step_{self.global_steps}')
            critic_remote_path = None if self.config.trainer.default_hdfs_dir is None else os.path.join(
                self.config.trainer.default_hdfs_dir, 'critic')
            self.critic_wg.save_checkpoint(critic_local_path, critic_remote_path)

        # save save_idx_reward
        if self.config.data.save_idx_reward:
            local_global_step_folder = os.path.join(self.config.trainer.default_local_dir, 'actor',
                                                f'global_step_{self.global_steps}')
            output_path = os.path.join(local_global_step_folder, f"modified_dataset.jsonl")
            with open(output_path, "w", encoding="utf-8") as f:
                for i, sample in enumerate(self.modified_dataset):
                    f.write(json.dumps(sample, ensure_ascii=False) + "\n")

    def _balance_batch(self, batch: DataProto, metrics, logging_prefix='global_seqlen'):
        """Reorder the data on single controller such that each dp rank gets similar total tokens"""
        attention_mask = batch.batch['attention_mask']
        batch_size = attention_mask.shape[0]
        global_seqlen_lst = batch.batch['attention_mask'].view(batch_size, -1).sum(-1).tolist()  # (train_batch_size,)
        world_size = self.actor_rollout_wg.world_size
        global_partition_lst = get_seqlen_balanced_partitions(global_seqlen_lst,
                                                              k_partitions=world_size,
                                                              equal_size=True)
        # reorder based on index. The data will be automatically equally partitioned by dispatch function
        global_idx = torch.tensor([j for partition in global_partition_lst for j in partition])
        batch.reorder(global_idx)
        global_balance_stats = log_seqlen_unbalance(seqlen_list=global_seqlen_lst,
                                                    partitions=global_partition_lst,
                                                    prefix=logging_prefix)
        metrics.update(global_balance_stats)

    def fit(self):
        """
        The training loop of PPO.
        The driver process only need to call the compute functions of the worker group through RPC to construct the PPO dataflow.
        The light-weight advantage computation is done on the driver process.
        """
        from verl.utils.tracking import Tracking
        from omegaconf import OmegaConf

        logger = Tracking(project_name=self.config.trainer.project_name,
                          experiment_name=self.config.trainer.experiment_name,
                          default_backend=self.config.trainer.logger,
                          config=OmegaConf.to_container(self.config, resolve=True))

        self.global_steps = 0

        # perform validation before training
        if self.val_reward_fn is not None and self.config.trainer.get('val_before_train', True):
            val_metrics = self._validate()
            pprint(f'Initial validation metrics: {val_metrics}')
            logger.log(data=val_metrics, step=self.global_steps)
            if self.config.trainer.get('val_only', False):
                return

        # we start from step 1
        self.global_steps += 1

        batch = None
        num_prompt_in_batch = 0
        num_gen_batches = 0

        for _ in range(self.config.trainer.total_epochs):
            
            for batch_dict in self.train_dataloader:
                new_batch: DataProto = DataProto.from_single_dict(batch_dict)

                metrics = {}
                timing_raw = {}

                # ====== Core modification: repeat first ======
                new_batch = new_batch.repeat(
                    repeat_times=self.config.actor_rollout_ref.rollout.n,
                    interleave=True
                )

                # # New: 10% chance to insert NO_THINK
                # new_batch = append_no_think_prompt(
                #     batch=new_batch,
                #     no_think_ids=NO_THINK_PROMPT_IDS,
                #     probability=0.1,
                #     pad_token_id=self.tokenizer.pad_token_id,
                #     # max_length=self.config.data.max_prompt_length
                # )

                # pop those keys for generation
                gen_batch = new_batch.pop(batch_keys=['input_ids', 'attention_mask', 'position_ids'])

                with _timer('step', timing_raw):
                    # generate a batch
                    with _timer('gen', timing_raw):
                        gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch)

                    # new_batch.non_tensor_batch['uid'] = np.array([str(uuid.uuid4()) for _ in range(len(new_batch.batch))],
                    #                                          dtype=object)
                    # ====== repeat AFTER ======
                    # new_batch = new_batch.repeat(
                    #     repeat_times=self.config.actor_rollout_ref.rollout.n,
                    #     interleave=True
                    # )
                    # ====== Correct UID allocation logic ======
                    # union之后
                    new_batch = new_batch.union(gen_batch_output)

                    # # Generate a UID based on the number of original prompts, and every n responses share one UID
                    bsz = len(new_batch.batch['input_ids']) // self.config.actor_rollout_ref.rollout.n
                    n = self.config.actor_rollout_ref.rollout.n

                    assert len(new_batch.batch['input_ids']) == bsz * n, "Batch size mismatch after generation."

                    uids = []
                    for _ in range(bsz):
                        shared_uid = str(uuid.uuid4())
                        uids.extend([shared_uid] * n)

                    # Insert new uid
                    new_batch.non_tensor_batch['uid'] = np.array(uids, dtype=object)


                    with _timer('reward', timing_raw):
                        # compute scores. Support both model and function-based.
                        # We first compute the scores using reward model. Then, we call reward_fn to combine
                        # the results from reward model and rule-based results.
                        if self.use_rm:
                            # we first compute reward model score
                            reward_tensor = self.rm_wg.compute_rm_score(new_batch)
                            new_batch = new_batch.union(reward_tensor)

                        # we combine with rule-based rm
                        reward_extra_infos_dict: dict[str, list]
                        try:
                            reward_result = self.reward_fn(new_batch, return_dict=True)
                            reward_tensor = reward_result['reward_tensor']
                            reward_extra_infos_dict = reward_result['reward_extra_info']
                        except Exception as e:
                            print(f'Error in reward_fn: {e}')
                            reward_tensor = self.reward_fn(new_batch)
                            reward_extra_infos_dict = {}

                        new_batch.batch['token_level_scores'] = reward_tensor

                        print(f'{list(reward_extra_infos_dict.keys())=}')
                        if reward_extra_infos_dict:
                            new_batch.non_tensor_batch.update({
                                k: np.array(v).astype(object) for k, v in reward_extra_infos_dict.items()
                            })

                        # compute rewards. apply_kl_penalty if available
                        if not self.config.actor_rollout_ref.actor.get('use_kl_loss', False):
                            new_batch, kl_metrics = apply_kl_penalty(new_batch,
                                                                     kl_ctrl=self.kl_ctrl,
                                                                     kl_penalty=self.config.algorithm.kl_penalty)
                            metrics.update(
                                kl_metrics)  # TODO: This will be cleared if we use multiple genenration batches
                        else:
                            new_batch.batch['token_level_rewards'] = new_batch.batch['token_level_scores']

                    with _timer('update_buffer_reward', timing_raw):
                        if "index" in new_batch.non_tensor_batch:
                            # print(">> [DEBUG] Starting reward update based on UID aggregation")

                            indexs = new_batch.non_tensor_batch["index"]
                            rewards = new_batch.batch["token_level_scores"].sum(dim=-1).cpu().numpy()
                            accs = new_batch.non_tensor_batch['is_correct_flags']
                            # Aggregate reward
                            idx_to_rewards = defaultdict(list)
                            idx_to_acc = defaultdict(list)
                            for idx, reward, acc in zip(indexs, rewards, accs):
                                idx_to_rewards[idx].append(reward)
                                idx_to_acc[idx].append(acc)

                            # Update the reward for each sample in modified_dataset
                            # updated_count = 0
                            for i, sample in enumerate(self.modified_dataset):
                                sample_idx = sample.get("index")
                                if sample_idx in idx_to_rewards:
                                    avg_reward = float(np.mean(idx_to_rewards[sample_idx]))
                                    avg_acc = float(np.mean(idx_to_acc[sample_idx]))
                                    sample["reward"] = avg_reward
                                    sample["acc"] = avg_acc
                                    
                    if not self.config.algorithm.filter_groups.enable:
                        batch = new_batch
                    else:  # NOTE: When prompts after filtering is less than train batch size, we skip to the next generation batch
                        metric_name = self.config.algorithm.filter_groups.metric
                        if metric_name == "seq_final_reward":
                            # Turn to numpy for easier filtering
                            new_batch.non_tensor_batch["seq_final_reward"] = new_batch.batch['token_level_scores'].sum(
                                dim=-1).numpy().astype(object)
                        elif metric_name == "seq_reward":
                            new_batch.non_tensor_batch["seq_reward"] = new_batch.batch['token_level_scores'].sum(
                                dim=-1).numpy().astype(object)

                        # Collect the sequence reward for each trajectory
                        prompt_uid2metric_vals = defaultdict(list)
                        for uid, metric_val in zip(new_batch.non_tensor_batch['uid'],
                                                   new_batch.non_tensor_batch[metric_name]):
                            prompt_uid2metric_vals[uid].append(metric_val)

                        prompt_uid2metric_std = {}
                        for prompt_uid, metric_vals in prompt_uid2metric_vals.items():
                            prompt_uid2metric_std[prompt_uid] = np.std(metric_vals)

                        kept_prompt_uids = [uid for uid, std in prompt_uid2metric_std.items() if std > 0]
                        num_prompt_in_batch += len(kept_prompt_uids)

                        kept_traj_idxs = []
                        for idx, traj_from_prompt_uid in enumerate(new_batch.non_tensor_batch['uid']):
                            if traj_from_prompt_uid in kept_prompt_uids:
                                kept_traj_idxs.append(idx)

                        new_batch = new_batch[kept_traj_idxs]
                        if batch is None:
                            batch = new_batch
                        else:
                            for k, v in new_batch.non_tensor_batch.items():
                                print(f"[{k}] type: {type(v)}, dtype: {getattr(v, 'dtype', 'N/A')}")
                                assert isinstance(v, np.ndarray), f"{k} is not a numpy array"
                                assert v.dtype == object, f"{k} is not dtype=object (got {v.dtype})"
                            batch = DataProto.concat([batch, new_batch])

                        prompt_bsz = self.config.data.train_batch_size
                        if num_prompt_in_batch < prompt_bsz:
                            print(f'{num_prompt_in_batch=} < {prompt_bsz=}')
                            max_num_gen_batches = self.config.algorithm.filter_groups.max_num_gen_batches
                            if max_num_gen_batches <= 0 or num_gen_batches < max_num_gen_batches:
                                print(f'{num_gen_batches=}. Keep generating...')
                                continue
                            else:
                                raise ValueError(
                                    f'{num_gen_batches=} >= {max_num_gen_batches=}. Generated too many. Please check your data.'
                                )
                        else:
                            # Align the batch
                            traj_bsz = self.config.data.train_batch_size * self.config.actor_rollout_ref.rollout.n
                            batch = batch[:traj_bsz]

                    # balance the number of valid tokens on each dp rank.
                    # Note that this breaks the order of data inside the batch.
                    # Please take care when you implement group based adv computation such as GRPO and rloo
                    if not isinstance(batch, DataProto): # Strange bug
                        batch = dataprotoitem_to_dataproto(batch)

                    self._balance_batch(batch, metrics=metrics)

                    # compute global_valid tokens
                    batch.meta_info['global_token_num'] = torch.sum(batch.batch['attention_mask'], dim=-1).tolist()
                    
                    # recompute old_log_probs
                    with _timer('old_log_prob', timing_raw):
                        old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)
                        batch = batch.union(old_log_prob)

                    if self.use_reference_policy:
                        # compute reference log_prob
                        with _timer('ref', timing_raw):
                            ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)
                            batch = batch.union(ref_log_prob)
                                
                    # compute values
                    if self.use_critic:
                        with _timer('values', timing_raw):
                            values = self.critic_wg.compute_values(batch)
                            batch = batch.union(values)

                    with _timer('adv', timing_raw):
                        batch.batch['token_level_rewards'] = batch.batch['token_level_scores']

                        # compute advantages, executed on the driver process
                        batch = compute_advantage(batch,
                                                  adv_estimator=self.config.algorithm.adv_estimator,
                                                  gamma=self.config.algorithm.gamma,
                                                  lam=self.config.algorithm.lam,
                                                  num_repeat=self.config.actor_rollout_ref.rollout.n)


                    # update critic
                    if self.use_critic:
                        with _timer('update_critic', timing_raw):
                            critic_output = self.critic_wg.update_critic(batch)
                        critic_output_metrics = reduce_metrics(critic_output.meta_info['metrics'])
                        metrics.update(critic_output_metrics)

                    # implement critic warmup
                    if self.config.trainer.critic_warmup <= self.global_steps:
                        # update actor
                        with _timer('update_actor', timing_raw):
                            actor_output = self.actor_rollout_wg.update_actor(batch)
                        actor_output_metrics = reduce_metrics(actor_output.meta_info['metrics'])
                        metrics.update(actor_output_metrics)

                    # validate
                    if self.val_reward_fn is not None and self.config.trainer.test_freq > 0 and \
                        self.global_steps % self.config.trainer.test_freq == 0:
                        with _timer('testing', timing_raw):
                            val_metrics: dict = self._validate()
                        metrics.update(val_metrics)

                    if self.config.trainer.save_freq > 0 and \
                            self.global_steps % self.config.trainer.save_freq == 0:
                        with _timer('save_checkpoint', timing_raw):
                            self._save_checkpoint()

                # collect metrics
                metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic, tokenizer=self.tokenizer))
                metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))

                metrics["train/num_gen_batches"] = num_gen_batches
                batch = None
                num_prompt_in_batch = 0
                num_gen_batches = 0
                
                # TODO: make a canonical logger that supports various backend
                logger.log(data=metrics, step=self.global_steps)

                self.global_steps += 1

                if self.global_steps >= self.total_training_steps:

                    # perform validation after training
                    if self.val_reward_fn is not None:
                        val_metrics = self._validate()
                        pprint(f'Final validation metrics: {val_metrics}')
                        logger.log(data=val_metrics, step=self.global_steps)
                    return
